#-*- coding:utf-8 -*-

import os,sys
import math
from opencv import cv,adaptors
from PIL import Image
import SimpleFilter

class FeatureVector:
    def __init__(self,image,maxsize=0):
        if maxsize==0:
            self.image=image
        else:
            self.image=SimpleFilter.Zoom(image,maxsize)


    def GetVector(self):
        if not hasattr(self,"_v"):
            self._calc()
        return self._v
    
    def _calc(self):
        raise NotImplementedError("implement _calc() which set the feature vector of self.image into self._v")

class CCV(FeatureVector):
    def __init__(self,image,th=4,useCIELab=False):
        FeatureVector.__init__(self,image)
        self.th=th
        self.__useLAB=useCIELab
        
    def _calc(self):
        #resize image
        img2=self.image
        #gaussian blur
        
        _img=adaptors.PIL2Ipl(img2)
        width,height=img2.size
        _img2=cv.cvCreateImage(cv.cvSize(width,height),cv.IPL_DEPTH_8U,3)
        cv.cvSmooth(_img,_img2,cv.CV_GAUSSIAN,5)
        img2=adaptors.Ipl2PIL(_img2)
        #cv.cvReleaseImage(_img)
        #cv.cvReleaseImage(_img2)
        del _img
        del _img2
        
        
        #reduce color into HTML16
        if self.__useLAB: SimpleFilter.ReduceColorLAB(img2,False)
        else: SimpleFilter.ReduceColor(img2,False)

        cnts=_CalcImageLabel(img2).GetLabelCounts()
        self._v=[0]*(len(SimpleFilter.colors)*2)
        
        for cnt in cnts:
            for i in xrange(len(SimpleFilter.colors)):
                if SimpleFilter.colors[i]==cnt[0]:
                    if cnt[1]>=self.th:
                        self._v[2*i]+=cnt[1]
                    else:
                        self._v[2*i+1]+=cnt[1]
        del img2

class _CalcImageLabel:
    def __init__(self,image):
        self.image=image
    
    def GetLabelCounts(self):
        if not hasattr(self,"__label"):
            self.__calc()
        return self.__labelcount

    def __calc(self):
        self.__label=dict()
        self.__labelcount=[]
        
        nextlabel = 1
        w,h=self.image.size
        sys.setrecursionlimit(w*h) #最悪の場合w*hすべてを再帰処理する

        pix=self.image.load()
        for x in xrange(w):
            for y in xrange(h):
                if not self.__label.has_key((x,y)):
                    cnt=[pix[x,y],0]
                    self.__label[(x,y)]=nextlabel+1
                    nextlabel+=1
                    self.__calc2(pix,x,y,cnt)
                    self.__labelcount.append(cnt)
        
                    
    
    def __calc2(self,pix,x,y,cnt):
        w,h=self.image.size
        if x<0 or x>=w: return
        if y<0 or y>=h: return
        
        cnt[1]+=1
        
        self.__calc3(pix,x+1,y-1,cnt,self.__label[(x,y)])
        self.__calc3(pix,x+1,y,cnt,self.__label[(x,y)])
        self.__calc3(pix,x+1,y+1,cnt,self.__label[(x,y)])
        self.__calc3(pix,x,y+1,cnt,self.__label[(x,y)])
        self.__calc3(pix,x-1,y+1,cnt,self.__label[(x,y)])
        
        
    def __calc3(self,pix,x,y,cnt,label):
        w,h=self.image.size
        if x<0 or x>=w: return
        if y<0 or y>=h: return
        
        if (not self.__label.has_key((x,y))) and pix[x,y]==cnt[0]:
            self.__label[(x,y)]=label
            self.__calc2(pix,x,y,cnt)

class HOG(FeatureVector):
    cellRow=4
    cellCol=4
    blockWidth=3
    blockHeight=3
    histgramDimension=9

    def _calc(self):
        img2=self.image.convert("L")
        width,height=img2.size
        pix=img2.load()
        
        #pix=dict()
        #for x in xrange(width):
            #for y in xrange(height):
                #pix[x,y]=sum(_pix[x,y])/3

        cellWidth=width/self.cellCol
        cellHeight=height/self.cellRow
        histgram=[0]*self.histgramDimension
        cellFeatureVector=[0]*self.histgramDimension*self.cellCol*self.cellRow

        for cy in xrange(self.cellRow):
            for cx in xrange(self.cellCol):
                for i in xrange(len(histgram)): histgram[i]=0
                for v in xrange(cellHeight):
                    y=cy*cellHeight+v
                    ymm=0 if y==0 else y-1
                    ypp=y if y==height-1 else y+1
                    for u in xrange(cellWidth):
                        x=cx*cellWidth+u
                        xmm=0 if x==0 else x-1
                        xpp=x if x==width-1 else x+1

                        dy=pix[x,ypp] - pix[x,ymm]
                        dx=pix[xpp,y] - pix[xmm,y]

                        idx=int(0.5*(math.atan2(dy,dx)/math.pi + 1)*(self.histgramDimension-1))
                        histgram[idx]+=math.sqrt(dx**2+dy**2)
                        
                for d in xrange(self.histgramDimension):
                    cellFeatureVector[(d*self.cellRow+cy)*self.cellCol+cx]=histgram[d]
        del histgram
        
        blockDimension=self.histgramDimension*self.blockWidth*self.blockHeight
        blockFeatureVector=[0]*blockDimension
        self._v=[0]*(self.histgramDimension*self.blockWidth*self.blockHeight*(self.cellRow-self.blockHeight)*(self.cellCol-self.blockWidth))

        for cy in xrange(self.cellRow-self.blockHeight):
            for cx in xrange(self.cellCol-self.blockWidth):
                for v in xrange(self.blockHeight):
                    for u in xrange(self.blockWidth):
                        for dd in xrange(self.histgramDimension):
                            blockFeatureVector[(dd*self.blockHeight+v)*self.blockWidth+u]=cellFeatureVector[(dd*self.cellRow+(cy+v))*self.cellCol+(cx+u)]
                norm=math.sqrt(sum([val*val for val in blockFeatureVector])+1)
                for v in xrange(blockDimension):
                    blockFeatureVector[v]/=norm
            
                for d in xrange(blockDimension):
                    self._v[(d*(self.cellRow-self.blockHeight)+cy)*(self.cellCol-self.blockWidth)+cx]=blockFeatureVector[d]
            
class GIST(FeatureVector):
    '''(非推奨) 怪しい & 死ぬほど遅い'''
    numBlock=4
    orientationPerScale=[8,8,4]
    
    def _calc(self):
        img2=self.image
        pix=img2.load()
        imgpix=dict()
        for x in xrange(img2.size[0]):
            for y in xrange(img2.size[1]):
                imgpix[x,y]=list(pix[x,y])
        self._v=self.__color_gist_scaletab(imgpix,img2.size,self.numBlock,self.orientationPerScale)

    @staticmethod
    def __color_gist_scaletab(img,size,numberBlocks,orientations):
        imageList = GIST._create_gabor(orientations,size)
       
        GIST._color_prefilt(img,size, 4)
        
        g = GIST._color_gist_gabor(img,size,numberBlocks,imageList)
        del imageList

        for descriptor in g:
            if not GIST._finite(descriptor):
                raise Exception("color_gist_scaletab() - descriptor is not valid! (NaN or infinite)")
        return g

    @staticmethod
    def _create_gabor(orientations,size):
        nFilters = sum(orientations)
        param=[[0.35]*nFilters,[0]*nFilters,[0]*nFilters,[0]*nFilters]
        cnt=0
        for i in xrange(len(orientations)):
            p1=0.3 / pow(1.85,i)
            p2=orientations[1]**2 / 64.0
            p3=math.pi/orientations[i]
            for j in xrange(orientations[i]):
                param[1][cnt]=p1
                param[2][cnt]=p2
                param[3][cnt]=p3*j
                cnt+=1

        width,height=size
        f=[[0]*width*height,[0]*width*height]
        cnt=0
        for y in xrange(height):
            for x in xrange(width):
                fx = x - width / 2.0
                fy = y - height / 2.0
                f[0][cnt]=math.sqrt(fx * fx + fy * fy)
                f[1][cnt]=math.atan2(fy, fx)
                cnt+=1
        
        GIST._fftshift(f[0], width, height)
        GIST._fftshift(f[1], width, height)
        
        imageList=[]
        for i in xrange(nFilters):
            fi=0
            img=dict()

            tmp2=-10.0*param[0][i]
            tmp3=(height*param[1][i],width*param[1][i])
            tmp4=2.0*param[2][i]*math.pi
            for iy in xrange(height):
                for ix in xrange(width):
                    tmp=f[1][fi]+param[3][i]
                    
                    if tmp<-math.pi:
                        tmp+=2*math.pi
                    elif tmp>math.pi:
                        tmp-=2*math.pi
                    img[ix, iy] = math.exp(tmp2*(f[0][fi]/tmp3[0] - 1)*(f[0][fi]/tmp3[1] - 1) - tmp4*tmp*tmp)
                    fi+=1
            imageList.append(img)
        del param
        del f

        return imageList
        
    @staticmethod
    def _fftshift(data,width,height):
        buff=list(data)

        for j in xrange((height+1)/2):
            for i in xrange((width+1)/2):
                data[(j+height/2)*width+i+width/2]=buff[j*width+i]

            for i in xrange(width/2):
                data[(j+height/2)*width+i] = buff[j*width+i+(width+1)/2]

        for j in xrange(height/2):
            for i in xrange((width+ 1)/2):
                data[j*width+i+width/2] = buff[(j+(height+1)/2)*width+i]
            for i in xrange(width/2):
                data[j*width+i] = buff[(j+(height+1)/2)*width + i + (width + 1) / 2]

        del buff
            
    @staticmethod
    def _color_prefilt(img,size,fc):
        padding = 5
        s1 = fc / math.sqrt(math.log(2))

        width,height=size
        width+=2*padding
        height+=2*padding
        
        src=dict()
        for y in xrange(size[1]):
            for x in xrange(size[0]):
                src[x + padding, y + padding]=[math.log(1.0+v) for v in img[x,y]]

            for x in xrange(padding):
                src[x, y + padding]=list(src[2 * padding - 1 - x, y + padding])
                src[x + padding + size[0], y + padding]=list(src[padding + size[0] - 1 - x, y + padding])

        for y in xrange(padding):
            for x in xrange(width):
                src[x, y]=list(src[x, 2 * padding - 1 - y])
                src[x, y + padding + size[1]]=list(src[x,padding + size[1] - 1 - y])
        
        gfc=[0]*width*height
        inData=[dict(),dict(),dict()]#[dict()]*3 はNG!!
        for y in xrange(height):
            for x in xrange(width):
                for i in xrange(3):
                    inData[i][x,y]=src[x, y][i]
                fx=x-width/2.0
                fy=y-height/2.0
                
                gfc[y*width+x]=math.exp(-(fx**2 + fy**2) / (s1**2))
        
        GIST._fftshift(gfc, width, height)

        #FFT
        outData=[Fourier.FFT2D(dt,width,height) for dt in inData]

        #Apply Whitening Filter
        for i in xrange(3):
            for y in xrange(height):
                for x in xrange(width):
                    outData[i][x,y]*=gfc[y * width + x]

        #IFFT
        inData2=[Fourier.IFFT2D(dt,width,height) for dt in outData]
        
        #Local Contrast Normalization
        for y in xrange(height):
            for x in xrange(width):
                for i in xrange(3):
                    src[x,y][i]-=inData2[i][x,y].real#/width*height
                mean=sum(src[x,y])/3
                
                inData[0][x,y]=mean**2

        #FFT
        outData[0]=Fourier.FFT2D(inData[0],width,height)
        
        #Apply Contrast Normalization Filter
        for y in xrange(height):
            for x in xrange(width):
                outData[0][x,y]*=gfc[y*width+x]

        #IFFT
        inData2[0] = Fourier.IFFT2D(outData[0], width, height)

        for y in xrange(height):
            for x in xrange(width):
                #val=math.sqrt(abs(inData2[0][x,y])/(width*height))
                val=math.sqrt(abs(inData2[0][x,y]))
                for i in xrange(3):
                    src[x,y][i]/=(0.2+val)

        del inData,inData2,outData,gfc

        for y in xrange(size[1]):
            for x in xrange(size[0]):
                for i in xrange(3):
                    img[x,y][i]=src[x+padding,y+padding][i]
        del src

    @staticmethod
    def _finite(num):
        if num==float("inf"): #inf
            return False
        elif num!=num: #NaN
            return False
        else:
            return True
    
    @staticmethod
    def _color_gist_gabor(img,size,w,G):
        inData=[dict(),dict(),dict()]#[dict()]*3
        for y in xrange(size[1]):
            for x in xrange(size[0]):
                for i in xrange(3):
                    inData[i][x,y] = img[x, y][i]
        
        #FFT
        outData=[Fourier.FFT2D(dt,size[0],size[1]) for dt in inData]
        outData2=[dict(),dict(),dict()]#[dict()]*3
        
        res=[0]*(3*w*w*len(G))
        #Filter
        for k in xrange(len(G)):
            for y in xrange(size[1]):
                for x in xrange(size[0]):
                    for i in xrange(3):
                        outData2[i][x, y] = outData[i][x, y] * G[k][x, y]

            #IFFT
            inData2=[Fourier.IFFT2D(dt,size[0],size[1]) for dt in outData2]
    
            for y in xrange(size[1]):
                for x in xrange(size[0]):
                    for i in xrange(3):
                        img[x, y][i] = abs(inData2[i][x, y])# / (size[0]*size[1])
    
            for i in xrange(3):
                GIST._color_down_N(res, i, len(G), k, img,size, w)
                
            del inData2
            
        del inData,outData,outData2

        return res

    @staticmethod
    def _color_down_N(res, c, nImage, K, src, size,N):
        offset=c * nImage * N * N + K * N * N
        nx=[i*size[0]/N for i in xrange(N+1)]
        ny=[i*size[1]/N for i in xrange(N+1)]

        for l in xrange(N):
            for k in xrange(N):
                s=sum((src[i,j][c] for i in xrange(nx[k],nx[k+1]) for j in xrange(ny[l],ny[l+1])))
                denom = (ny[l+1] - ny[l]) * (nx[k+1] - nx[k])
                res[offset + k*N + l] = s/denom   